# -*- coding:UTF-8 -*-
from django.conf import settings
from celery import task
import time
import os
import sys
import zipfile
#from upload.result.test_crnn_dfa import ImgPreProcess
import numpy as np
import cv2
# crnn packages
import torch
from torch.autograd import Variable
from .result import myutils as utils
from .result.models import crnn as crnn
from .result import alphabets as alphabets
from .result import params as params
from .result.dfa import SWFilter
import imutils
import glob, shutil,traceback
sys.path.append(os.getcwd())
# ---params
# (1)url the path of images in apk
# ---return 
# (1)the number of hramful messages
# (2)the list of harmful pictures' path
# self.update_state(state='PROGRESS',meta={'i':i})
# python36 manage.py celery worker -c 4 --loglevel=info


def unzip(filename: str):
    try:
        file = zipfile.ZipFile(filename)
        dirname = filename.replace('.zip', '')
        # 如果存在与压缩包同名文件夹 提示信息并跳过
        if os.path.exists(dirname):
            tmp_path = os.getcwd()
            os.chdir(dirname)
            for f in os.listdir(dirname):
                if os.path.isfile(f):
                    os.remove(f)
                else:
                    shutil.rmtree(f)
            os.chdir(tmp_path)
            file.extractall(dirname)
            file.close()
            # 递归修复编码
            rename(dirname)
        else:
            # 创建文件夹，并解压
            os.mkdir(dirname)
            file.extractall(dirname)
            file.close()
            # 递归修复编码
            rename(dirname)
    except Exception as e:
        print(f'{filename} unzip fail error is:' ,e)
        traceback.print_exc()


def rename(pwd: str, filename=''):
    """压缩包内部文件有中文名, 解压后出现乱码，进行恢复"""
    
    path = f'{pwd}/{filename}'
    if os.path.isdir(path):
        for i in os.scandir(path):
            rename(path, i.name)
    newname = filename.encode('cp437').decode('gbk')
    os.rename(path, f'{pwd}/{newname}')

total_num = 0
def ListFilesToTxt(dir,file,wildcard,recursion):
    exts = wildcard.split(" ")
    files = os.listdir(dir)
    for name in files:
        fullname=os.path.join(dir,name)
        if(os.path.isdir(fullname) & recursion):
            ListFilesToTxt(fullname,file,wildcard,recursion)
        else:
            for ext in exts:
                if(name.endswith(ext)):
                    fullname = fullname.replace('/var/www/Imgdetection/upload/images/','')
                    global total_num
                    total_num = total_num + 1
                    file.write(fullname + "\n")
                    break

def Test(dir,apkname):
    #outfile = dir + "/imgpath.txt"                     #写入的txt文件名
    outfile =  os.path.join(settings.MEDIA_ROOT,'result',apkname.split('.')[0]+"_res.txt") 
    wildcard = ".jpg .png .jpeg .bmp .tif"      #要读取的文件类型；
    file = open(outfile,'w+')
    if not file:
        print ("cannot open the file %s for writing" % outfile)
    ListFilesToTxt(dir,file,wildcard, 1)
    global total_num
    total_num_img = total_num
    total_num = 0
    file.close()
    return outfile,total_num_img

# def Testimage(res_name):
#     #res_name='2.apk'
#     apkurl = os.path.join(settings.MEDIA_ROOT,'images',res_name) 
#     zipurl = os.path.join(settings.MEDIA_ROOT,'images',res_name.split('.')[0]+'.zip')
#     os.rename(apkurl,zipurl)
    
#     unzip(zipurl)
#     url = os.path.join(settings.MEDIA_ROOT,'images',res_name.split('.')[0])
    
#     outurl,file = Test(url,res_name)
#     return 9,[]
#     flag,hlist = ImgPreProcess(outurl)
#     return flag,hlist 
#     #with open(url+'/imgpath.txt', 'r', encoding='utf-8') as file:
#     #  for img_path in file.readlines() :
#     #        img_path=img_path.strip('\n')
#     #        test2.extract(img_path)
#     #   file.close()

#######################################################################################

# crnn params
str1 = alphabets.alphabet
crnn_model_path = '/var/www/Imgdetection/upload/result/trained_models/crnn_Rec_done_0_0.9234375.pth'
alphabet = str1
nclass = len(alphabet) + 1

# crnn文本信息识别
def crnn_recognition(cropped_image, model):
    converter = utils.strLabelConverter(alphabet)
    try:
        image = cv2.cvtColor(cropped_image, cv2.COLOR_BGR2GRAY)
    except:
        return 0
    else:
        ### ratio
        ### 280是中文训练集中图片的宽度，160是将原始图片缩小后的图片宽度
        w_now = int(image.shape[1] / (500 * 1.0 / params.imgW))
        h, w = image.shape
        image = cv2.resize(image, (0,0), fx=w_now/w, fy=params.imgH/h, interpolation=cv2.INTER_CUBIC)
        image = (np.reshape(image, (params.imgH, w_now, 1))).transpose(2, 0, 1)
        image = image.astype(np.float32) / 255.
        image = torch.from_numpy(image).type(torch.FloatTensor)
        image.sub_(params.mean).div_(params.std)
        if torch.cuda.is_available():
            image = image.cuda()
        image = image.view(1, *image.size())
        image = Variable(image)

        model.eval()
        preds = model(image)

        _, preds = preds.max(2)
        preds = preds.transpose(1, 0).contiguous().view(-1)

        preds_size = Variable(torch.IntTensor([preds.size(0)]))
        sim_pred = converter.decode(preds.data, preds_size.data, raw=False)
        SWResult = SWFilter(format(sim_pred))
        #print('results: {0} '.format(sim_pred))   #打印识别结果
        #print('SWResult: {:30s} SWTime: {:10s}' .format(SWResult['Result'], SWResult['ElapsedT']))
        #打印过滤结果，及过滤时间
        return SWResult

def dfaFilter(pre_txt_url):
    _FLAG = 0
    # crnn network
    
    model = crnn.CRNN(32, 1, nclass, 256)
    if torch.cuda.is_available():
        model = model.cuda()
    #print('loading pretrained model from {0}'.format(crnn_model_path))
    # 导入已经训练好的crnn模型

    model.load_state_dict(torch.load(crnn_model_path, map_location=torch.device('cpu')))
    
    started = time.time()
    ## read an image
    #待检测图片文件夹路径
    imagepath = pre_txt_url+'/'
    #imagepath = 'test_images/Result/'
    _LIST = []
    _OUTFILE =  os.path.join(settings.MEDIA_ROOT,'result', str(int(time.time())) + "_res.txt") 
    with open(_OUTFILE,'w+') as f:
        for img in os.listdir(imagepath):
            #print(img)#part of the picture in apk
            image = cv2.imread(imagepath+img)
            #image = cv2.imread(opt.images_path)
            _FLAG = crnn_recognition(image, model)
            if _FLAG == 1:  #harmful part
                f.write(img+'\n')
                _LIST.append(imagepath+img)
            finished = time.time()
            #print('elapsed time: {0}'.format(finished-started))
    #检测全部结束，返回值是否存有有害信息，>=1为有，0为无
    #print(_FLAG)
    #print(_LIST)
    return _FLAG,_LIST,_OUTFILE
    
def ImgPreProcess(txt_url):
    # 创建输出目录
    out_dir = '/var/www/Imgdetection/upload/result/Result/'
    for fil in glob.glob(out_dir + '/*'):
         os.remove(fil)
    l = 1
    #有bug202008162323放弃，其余流程都已做好,敏感词库还未上传 
    #需要把子目录更改所有者，读写权限，即使没有子目录或子文件
    # txt_url = '/var/www/Imgdetection/upload/result/Testapk_res.txt'
    with open(txt_url,"r") as myf:              # 返回一个文件对象
        for eachline in myf:
            eachline = os.path.join('/var/www/Imgdetection/upload/images',eachline.strip('\n'))
            #eachline = eachline.strip('\n')
            # 读取每一张图片
            image = cv2.imread(eachline)
            if image is not None:
                # 转化为灰度图
                
                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
                
                # 形态计算处理，得到可以查找矩形的图片
                sobel = cv2.Sobel(gray, cv2.CV_8U, 1, 0, ksize=3)
                # 阈值处理，二值化
                
                ret, binary = cv2.threshold(sobel, 0, 255,
                                            cv2.THRESH_BINARY | cv2.THRESH_OTSU)
                
                # 设置膨胀和腐蚀操作的核函数
                elem1 = cv2.getStructuringElement(cv2.MORPH_RECT, (30, 9))
                elem2 = cv2.getStructuringElement(cv2.MORPH_RECT, (24, 6))
                # 膨胀一次，使得轮廓更加突出
                dilation = cv2.dilate(binary, elem2, iterations=1)
                # 腐蚀一次，去掉噪声细节，如表格竖线等
                erosion = cv2.erode(dilation, elem1, iterations=1)
                # 连续膨胀3次，使轮廓更明显
                dilation2 = cv2.dilate(erosion, elem2, iterations=1)

                # 查找和筛选文字区域
                region = []
                # 查找得到多个轮廓区域，存储在列表contours中
                contours, hierarchy = cv2.findContours(dilation2,cv2.RETR_TREE,
                                                    cv2.CHAIN_APPROX_SIMPLE)

            
                
                # 遍历所有轮廓，筛选面积较小的
                for i in range(len(contours)):
                    cnt = contours[i]
                    area = cv2.contourArea(cnt)
                    if(area < 800):
                        continue
                    # 将轮廓形状近似到另外一种由更少点组成的轮廓形状
                    epsilon = 0.001 * cv2.arcLength(cnt, True)
                    approx = cv2.approxPolyDP(cnt, epsilon,True)
                    # 找到最小矩形轮廓，及其坐标
                    rect = cv2.minAreaRect(cnt)
                    box = cv2.boxPoints(rect)
                    box = np.int0(box)
                    height = abs(box[0][1]-box[2][1])
                    width = abs(box[0][0]-box[2][0])

                    if(height > width * 1.3):
                        continue
                    region.append(box)
                    output = image
                    for box in region:
                        cv2.drawContours(output, [box], 0, (0, 0, 255), 1)
                        #print(box)
                    #cv2.imshow('Output', output)
                    #cv2.waitKey(0)
                        
                    Xs = [i[0] for i in box]
                    Ys = [i[1] for i in box]
                    x1 = min(Xs)
                    x2 = max(Xs)
                    y1 = min(Ys)
                    y2 = max(Ys)
                    height = y2 - y1
                    width = x2 - x1
                    crop_img:object = output[y1:y1+height, x1:x1+width]
                    #im = cv2.imshow('Crop_img', crop_img)
                    #cv2.waitKey(0)
                    try:
                        img_gray = cv2.cvtColor(crop_img, cv2.COLOR_BGR2GRAY)
                    except:
                        
                        continue
                    else:
                        # define a threshold, 128 is the middle of black and white in grey scale
                        thresh = 128
                                
                        # assign blue channel to zeros
                        img_binary = cv2.threshold(img_gray, thresh, 255, cv2.THRESH_BINARY)[1]
                                
                        #crop_img_gray=cv2.cvtColor(crop_img, cv2.COLOR_BGR2GRAY)
                        #print(out_dir+'/%d%d' %(l,i) +'.jpg')
                        size = (280,32)
                        img_binary = cv2.resize(img_binary,size)
                        try:
                            cv2.imwrite(out_dir+'%d%d' %(l,i)+'gray.jpg', img_binary)
                        except:
                            print('write img error')
            l = l + 1
    return out_dir


@task(bind=True)
def looploop(self,res_name):
    name = res_name.split('.')[-2]
    apkurl = os.path.join(settings.MEDIA_ROOT,'images',res_name) 
    zipurl = os.path.join(settings.MEDIA_ROOT,'images',name+'.zip')
    if res_name.split('.')[-1] != 'zip':
        if os.path.isfile(zipurl):
            os.remove(zipurl)
        os.rename(apkurl,zipurl)
    self.update_state(state='PROGRESS',meta={'i':5})
    unzip(zipurl)
    self.update_state(state='PROGRESS',meta={'i':10})
    url = os.path.join(settings.MEDIA_ROOT,'images',res_name.split('.')[0])    
    outurl,_TOTAL= Test(url,res_name)
    self.update_state(state='PROGRESS',meta={'i':15})
    out_dir = ImgPreProcess(outurl)
    self.update_state(state='PROGRESS',meta={'i':20})
    progressNum = 20
    #flag,hlist = dfaFilter(out_dir)
    pre_txt_url = out_dir
    _FLAG = 0
    # crnn network
    model = crnn.CRNN(32, 1, nclass, 256)
    if torch.cuda.is_available():
        model = model.cuda()
    #print('loading pretrained model from {0}'.format(crnn_model_path))
    # 导入已经训练好的crnn模型
    model.load_state_dict(torch.load(crnn_model_path, map_location=torch.device('cpu')))
    started = time.time()
    ## read an image
    #待检测图片文件夹路径
    imagepath = pre_txt_url+'/'
    #imagepath = 'test_images/Result/'
    _LIST = []
    _OUTFILE =  os.path.join(settings.MEDIA_ROOT,'result', str(int(time.time())) + "_res.txt") 
    data_length = len(os.listdir(imagepath))
    with open(_OUTFILE,'w+') as f:
        for i,img in enumerate(os.listdir(imagepath)):
            #print(img)#part of the picture in apk
            image = cv2.imread(imagepath+img)
            #image = cv2.imread(opt.images_path)
            temp_flag = crnn_recognition(image, model)
            _FLAG = _FLAG + temp_flag
            if temp_flag == 1:  #harmful part
                f.write(img+'\n')
                _LIST.append(img)
            finished = time.time()
            self.update_state(state='PROGRESS',meta={'i':progressNum + (i + 1) / data_length * 80})
            #print('elapsed time: {0}'.format(finished-started))
    #检测全部结束，返回值是否存有有害信息，>=1为有，0为无
    #print(_FLAG)
    #print(_LIST)
    return _FLAG,_LIST,_TOTAL,_OUTFILE
